from typing import TypedDict, Literal

from langchain_core.output_parsers import StrOutputParser
from langgraph.constants import START, END
from langgraph.graph import StateGraph
from pydantic import BaseModel, Field

from agent.my_llm import llm


class State(TypedDict):
    joke: str  # 生成的冷笑话内容
    topic: str  # 用户指定的主题
    feedback: str  # 改进建议
    funny_or_not: str  # 幽默评级


# 节点函数
def generator_func(state: State):
    """由大模型生成一个冷笑话的节点"""
    prompt = (
        f"根据反馈改进笑话：{state['feedback']}\n主题：{state['topic']}"
        if state.get("feedback", None)
        else f"创作一个关于{state['topic']}的笑话"
    )
    # 方式一：
    # resp = llm.invoke(prompt)
    # return {'joke': resp.content}
    # 方式二：
    chain = llm | StrOutputParser()
    resp = chain.invoke(prompt)
    return {'joke': resp}


# 结构化输出模型（用于LLM评估反馈）
class Feedback(BaseModel):
    """使用此工具来结构化你的响应"""
    grade: Literal["funny", "not funny"] = Field(
        description="判断笑话是否幽默",
        examples=["funny", "not funny"]
    )
    feedback: str = Field(
        description="若不幽默，提供改进建议",
        examples=["可以加入双关语或意外结局"]
    )


# 节点函数
def evaluator_func(state: State):
    """评估状态中的冷笑话"""
    # 方式一：
    chain = llm.with_structured_output(Feedback)
    resp = chain.invoke(
        f"评估此笑话的幽默程度：\n{state['joke']}\n"
        "注意：幽默应包含意外性或巧妙措辞"
    )
    return {
        'funny_or_not': resp.grade,
        'feedback': resp.feedback,
    }
    # 方式二：
    # chain = llm.bind_tools([Feedback])
    # evaluation = chain.invoke(
    #     f"评估此笑话的幽默程度：\n{state['joke']}\n"
    #     "注意：幽默应包含意外性或巧妙措辞"
    # )
    # evaluation = evaluation.tool_calls[-1]['args']
    # return {
    #     "funny_or_not": evaluation['grade'],
    #     "feedback": evaluation['feedback']
    # }


# 条件边的路由函数
def route_func(state: State) -> str:
    """动态路由决策函数"""
    return 'Accepted' if state.get('funny_or_not', None) == "funny" else "Rejected + Feedback"


builder = StateGraph(State)

# 节点定义
builder.add_node('generator', generator_func)
builder.add_node('evaluator', evaluator_func)

# 边定义
builder.add_edge(START, 'generator')
builder.add_edge('generator', 'evaluator')
builder.add_conditional_edges(
    'evaluator',
    route_func,
    {
        'Accepted': END,
        'Rejected + Feedback': 'generator'
    }
)

# 图构建
graph = builder.compile()

